#ifndef KNOWHERE_ASCEND_UTILS
#define KNOWHERE_ASCEND_UTILS

#include <mutex>
#include <thread>
#include <unordered_map>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <atomic>

#include "ascend_exception.h"
#include "knowhere/expected.h"
#include "knowhere/log.h"

namespace knowhere {
namespace ascend {
class ascend_device_manager {
public:
    static ascend_device_manager& instance();
    int random_choose() const;
    int choose_with_load(size_t load);
    void release_load(int32_t device_id, size_t load);
    std::mutex* get_dev_mutex_ptr(int32_t device_id);

private:
    ascend_device_manager();
    void init();
    std::vector<size_t> memory_load_;
    mutable std::mutex mtx_;
    std::atomic<bool> inited_{false};
    std::vector<std::mutex*> device_mutex_ptres;
};
} // namespace ascend
}

#define RANDOM_CHOOSE_ASCEND_DEVICE_WITH_ASSIGN(x)                             \
    do {                                                                \
        x = knowhere::ascend::ascend_device_manager::instance().random_choose(); \
    } while (0)
#define MIN_LOAD_CHOOSE_ASCEND_DEVICE_WITH_ASSIGN(x, load)                            \
    do {                                                                       \
        x = knowhere::ascend::ascend_device_manager::instance().choose_with_load(load); \
    } while (0)
#define RELEASE_ASCEND_DEVICE(x, load)                                           \
    do {                                                                  \
        knowhere::ascend::ascend_device_manager::instance().release_load(x, load); \
    } while (0)

#endif